# -*- coding: utf-8 -*-
"""
Created on Mon Aug 16 13:14:34 2021

@author: weixifei
"""
import numpy as np
from sklearn.metrics import cohen_kappa_score
import torch
import os
# from tqdm import *
from tqdm import tqdm
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def train(model, num_epochs, train_dataloader, val_dataloader, optimizer, criterion):
    train_loss_all=[]
    val_loss_all=[]
    scheduler = ReduceLROnPlateau(optimizer, 'min',factor=0.5, patience=5, verbose=True)
    for epoch in tqdm(range(num_epochs)):
        model.train()
        i=0
        loss_total=0.0
        avg_loss_list = []
        avg_kappa_list = []
        for i, (fun_images,map_images, labels) in enumerate(train_dataloader):
            map_images = map_images.to(device)
            fun_images = fun_images.to(device)
            labels = labels.to(device)

            logits = model(fun_images , map_images)
            loss = criterion(logits, labels)
            print(loss)
            for p, l in zip(logits.detach().cpu().numpy().argmax(1), labels.detach().cpu().numpy()):
                avg_kappa_list.append([p, l])

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_total+=loss

            avg_loss_list.append(loss.item())
            i+=1

#        scheduler.step(loss_total)
        train_loss = np.array(avg_loss_list).mean()
        avg_kappa_list = np.array(avg_kappa_list)
        train_loss_all.append(train_loss)
        train_kappa = cohen_kappa_score(avg_kappa_list[:, 0], avg_kappa_list[:, 1], weights='quadratic')
        avg_loss_list = []
        avg_kappa_list = []
        val_loss, val_kappa = val(model, val_dataloader, criterion)
        val_loss_all.append(val_loss)
        scheduler.step(val_loss)
        print("train:epochs={}/{} avg_loss={:.4f} train_kappa={:.4f}".format(epoch, num_epochs, train_loss, train_kappa))
        print("val: epochs={}/{} avg_loss={:.4f} val kappa={:.4f}".format(epoch, num_epochs, val_loss, val_kappa))
        if epoch >= 20:
            # best_kappa = avg_kappa
            torch.save(model.state_dict(),
                    os.path.join("../model_save_GCC1000/model{}_{:.4f}.pth".format(epoch,val_kappa)))

        # if (epoch+1)%20==0:
        #     plt.plot(train_loss_all)
        #     plt.plot(val_loss_all)
        #     # plt.show()
        #     plt.pause(1)


def val(model, val_dataloader, criterion):
    model.eval()
    avg_loss_list = []
    cache = []
    with torch.no_grad():
        for i, (fun_images,map_images, labels) in enumerate(val_dataloader):
            map_images = map_images.to(device)
            fun_images = fun_images.to(device)
            labels = labels.to(device)
            labels=labels.to(device)
            logits = model(fun_images, map_images)
            for p, l in zip(logits.detach().cpu().numpy().argmax(1), labels.detach().cpu().numpy()):
                cache.append([p, l])

            loss = criterion(logits, labels)
            # acc = paddle.metric.accuracy(input=logits, label=labels.reshape((-1, 1)), k=1)
            avg_loss_list.append(loss.item())
        cache = np.array(cache)
        kappa = cohen_kappa_score(cache[:, 0], cache[:, 1], weights='quadratic')
        avg_loss = np.array(avg_loss_list).mean()

    return avg_loss, kappa

if __name__ =="__main__":
    import torch
    print(torch.__version__)
    print(torch.version.cuda)
    print(torch.backends.cudnn.version())

